[ML] Refactor inference API service tests base classes #135461

jonathan-buttner · 2025-09-25T17:18:15Z

This PR refactors a base test class for testing service classes (e.g. OpenAiService).

Highlights

OpenAiServiceTests now leverages the base class
The parsePersistedConfig from OpenAI were moved to AbstractInferenceServiceParameterizedTests so they can be leveraged for all classes that extend from it
Created AbstractInferenceServiceParameterizedTests which uses parameterized tests to represent the parsePersistedConfig and parsePersistedConfigWithSecrets tests they can be removed from each subclassed tests
Moved parsePersistedConfigWithSecrets tests from the AbstractInferenceServiceTests to AbstractInferenceServiceParameterizedTests
Fixed a bug in the CustomService where we were parsing the chunking settings using the old default settings instead of the new ones

jonathan-buttner · 2025-09-25T17:21:03Z

...inference/src/main/java/org/elasticsearch/xpack/inference/services/custom/CustomService.java

-            var chunkingSettings = extractChunkingSettings(config, taskType);
+            ChunkingSettings chunkingSettings = null;
+            if (TaskType.TEXT_EMBEDDING.equals(taskType)) {
+                chunkingSettings = ChunkingSettingsBuilder.fromMap(removeFromMapOrDefaultEmpty(config, ModelConfigurations.CHUNKING_SETTINGS));


This is the bug fix. extractChunkingSettings was using removeFromMap which if the settings don't exist it would provide null to the ChunkingSettingsBuilder.fromMap(). We intentionally do this when parsing from persistent state to handle backwards compatibility I think.

I don't think the change here in parseRequestConfig will cause any backwards compatibility issues. For new endpoints being created we'll use the newer default chunking settings instead though.

jonathan-buttner · 2025-09-25T17:21:38Z

...icsearch/xpack/inference/services/openai/completion/OpenAiChatCompletionServiceSettings.java

    // To find this information you need to access your account's limits https://platform.openai.com/account/limits
    // 500 requests per minute
-    private static final RateLimitSettings DEFAULT_RATE_LIMIT_SETTINGS = new RateLimitSettings(500);
+    public static final RateLimitSettings DEFAULT_RATE_LIMIT_SETTINGS = new RateLimitSettings(500);


Making various things public so they can be accessible in tests that are outside of the package.

jonathan-buttner · 2025-09-25T17:22:15Z

.../plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/ServiceUtils.java

        );
    }

+    public static String parsePersistedConfigErrorMsg(String inferenceEntityId, String serviceName, TaskType taskType) {


Adding a new version of the error message to make the error more clear. Ideally all of the services will switch to use this one.

Would it be possible to switch the services to use the new method in this PR? There are a dozen tests that would need to be updated to reflect the new message, but it would be nice to have consistency across all services.

Yep I can do that 👍

jonathan-buttner · 2025-09-25T17:23:23Z

...ence/src/test/java/org/elasticsearch/xpack/inference/services/custom/CustomServiceTests.java

+                }
+
+                @Override
+                protected void assertRerankerWindowSize(RerankingInferenceService rerankingInferenceService) {


Adding this to the configuration of the tests so that it can be leveraged by CustomServiceTests and CustomServiceParameterizedTests.

jonathan-buttner · 2025-09-25T17:23:56Z

...ence/src/test/java/org/elasticsearch/xpack/inference/services/custom/CustomServiceTests.java

    }
-
-    @Override
-    public InferenceService createInferenceService() {


This was pushed into the abstract base class

jonathan-buttner · 2025-09-25T17:29:48Z

...a/org/elasticsearch/xpack/inference/services/AbstractInferenceServiceParameterizedTests.java

@@ -0,0 +1,387 @@
+/*


The reason this class exists and isn't included in AbstractInferenceServiceTests is adding parameterized tests in AbstractInferenceServiceTests results in the the non parameterized tests being run for each permutation, even though those tests don't depend on the parameterized test information.

It seemed cleaner to separate them 🤷‍♂️

jonathan-buttner · 2025-09-25T17:30:36Z

...a/org/elasticsearch/xpack/inference/services/AbstractInferenceServiceParameterizedTests.java

+                {
+                    new TestCaseBuilder(
+                        "Test parsing persisted config without chunking settings",
+                        testConfiguration -> getPersistedConfigMap(


The test configurations have to be passed in because the parameters() method must be static

jonathan-buttner · 2025-09-25T17:31:45Z

.../src/test/java/org/elasticsearch/xpack/inference/services/AbstractInferenceServiceTests.java

 */
-public abstract class AbstractInferenceServiceTests extends InferenceServiceTestCase {
-
-    protected final MockWebServer webServer = new MockWebServer();


Pulled this logic into the base class

jonathan-buttner · 2025-09-25T17:33:57Z

.../src/test/java/org/elasticsearch/xpack/inference/services/AbstractInferenceServiceTests.java

-        protected abstract Map<String, Object> createTaskSettingsMap();
-
-        protected abstract Map<String, Object> createSecretSettingsMap();
+    public void testParseRequestConfig_CreatesAnEmbeddingsModel() throws Exception {


The *ParseRequestConfig_* tests could be converted into parameterized tests as well. I decided to not do it in this PR because it's already quite a few changes.

jonathan-buttner · 2025-09-25T17:34:36Z

.../src/test/java/org/elasticsearch/xpack/inference/services/AbstractInferenceServiceTests.java

-        protected abstract Map<String, Object> createTaskSettingsMap();
-
-        protected abstract Map<String, Object> createSecretSettingsMap();
+    public void testParseRequestConfig_CreatesAnEmbeddingsModel() throws Exception {


I'm adding a few tests that were in OpenAiServiceTests to make them common.

…tor-tests-abstract

…ner/elasticsearch into ml-refactor-tests-abstract

elasticsearchmachine · 2025-09-25T19:26:39Z

Pinging @elastic/ml-core (Team:ML)

DonalEvans · 2025-09-26T18:38:49Z

.../plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/ServiceUtils.java

        );
    }

+    public static String parsePersistedConfigErrorMsg(String inferenceEntityId, String serviceName, TaskType taskType) {


Would it be possible to switch the services to use the new method in this PR? There are a dozen tests that would need to be updated to reflect the new message, but it would be nice to have consistency across all services.

DonalEvans · 2025-09-26T18:40:10Z

.../plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/ServiceUtils.java


+    public static String parsePersistedConfigErrorMsg(String inferenceEntityId, String serviceName, TaskType taskType) {
+        return format(
+            "Failed to parse stored model [%s] for [%s] service, error: [%s]. Please delete and add the service again",


Would deleting and adding the service again actually help if the task type was unsupported?

Yeah I think deleting is probably the only solution here. I think this would only occur if the the inference endpoint got corrupted somehow. Basically this is saying that the persisted inference endpoint is stating it is leverage a particular task type that is not supported. The request context parsing should prevent getting into that scenario. But if we had a regression or the endpoint was corrupted somehow we could.

DonalEvans · 2025-09-26T18:47:26Z

...gin/inference/src/main/java/org/elasticsearch/xpack/inference/services/ai21/Ai21Service.java

            serviceSettingsMap,
            secretSettingsMap,
-            parsePersistedConfigErrorMsg(modelId, NAME)
+            parsePersistedConfigErrorMsg(modelId, NAME, taskType)


To reduce some code duplication and prevent us from constructing a String that we might never use every time we call parseRequestConfig(), I think it should be possible to move the parsePersistedConfigErrorMsg() call down into createModel() along with the unsupportedTaskTypeErrorMsg() call used for the REQUEST context and only call them (which one we call would depend on the context passed in to createModel()) when throwing an exception that would need the message.

I haven't checked every Service to see if this would work for all of them, but I assume that the structure is pretty similar between implementations.

DonalEvans · 2025-09-26T18:53:32Z

...inference/src/main/java/org/elasticsearch/xpack/inference/services/custom/CustomService.java

-    private static ChunkingSettings extractChunkingSettings(Map<String, Object> config, TaskType taskType) {
-        if (TaskType.TEXT_EMBEDDING.equals(taskType)) {
-            return ChunkingSettingsBuilder.fromMap(removeFromMap(config, ModelConfigurations.CHUNKING_SETTINGS));
-        }
-
-        return null;
-    }


Would it be better to keep using this method in the two places below, where the behaviour is unchanged, rather than duplicating the logic?

…tor-tests-abstract

DonalEvans · 2025-09-30T14:25:01Z

...inference/src/main/java/org/elasticsearch/xpack/inference/services/custom/CustomService.java


+    private static ChunkingSettings extractPersistentChunkingSettings(Map<String, Object> config, TaskType taskType) {
+        if (TaskType.TEXT_EMBEDDING.equals(taskType)) {
+            // note there's


Unfinished comment

…ner/elasticsearch into ml-refactor-tests-abstract

…tor-tests-abstract

jonathan-buttner added 3 commits September 24, 2025 09:31

Refactoring openai

028d445

Splitting up parameterized tests

c3243fd

Working tests

399a50b

jonathan-buttner added >refactoring :ml Machine learning Team:ML Meta label for the ML team v9.2.0 labels Sep 25, 2025

elasticsearchmachine added 2 commits September 25, 2025 17:25

[CI] Auto commit changes from spotless

060eed9

[CI] Update transport version definitions

ac44cc6

jonathan-buttner commented Sep 25, 2025

View reviewed changes

jonathan-buttner added 2 commits September 25, 2025 13:35

Merge branch 'main' of github.com:elastic/elasticsearch into ml-refac…

0af841d

…tor-tests-abstract

Merge branch 'ml-refactor-tests-abstract' of github.com:jonathan-butt…

bbd0707

…ner/elasticsearch into ml-refactor-tests-abstract

jonathan-buttner marked this pull request as ready for review September 25, 2025 19:26

DonalEvans reviewed Sep 26, 2025

View reviewed changes

jonathan-buttner and others added 6 commits September 29, 2025 11:52

Merge branch 'main' of github.com:elastic/elasticsearch into ml-refac…

9392ec0

…tor-tests-abstract

Removing deprecated function

cc99292

Moving string creation and refactoring customservice chunking

0896e5c

Removing usages of persistent function

49c1301

Merge branch 'main' of github.com:elastic/elasticsearch into ml-refac…

3110967

…tor-tests-abstract

[CI] Auto commit changes from spotless

6ec99f9

DonalEvans reviewed Sep 30, 2025

View reviewed changes

jonathan-buttner added 3 commits September 30, 2025 10:35

Finishing comment

d52d4dc

Merge branch 'ml-refactor-tests-abstract' of github.com:jonathan-butt…

bbd1f3e

…ner/elasticsearch into ml-refactor-tests-abstract

Merge branch 'main' of github.com:elastic/elasticsearch into ml-refac…

dd9dde0

…tor-tests-abstract

DonalEvans approved these changes Sep 30, 2025

View reviewed changes

jonathan-buttner enabled auto-merge (squash) September 30, 2025 15:30

jonathan-buttner merged commit 698a7b6 into elastic:main Sep 30, 2025
34 checks passed

jonathan-buttner deleted the ml-refactor-tests-abstract branch September 30, 2025 15:50

[ML] Refactor inference API service tests base classes #135461

[ML] Refactor inference API service tests base classes #135461

Uh oh!

Conversation

jonathan-buttner commented Sep 25, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

elasticsearchmachine commented Sep 25, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants